import os, random, struct, hashlib
from base64 import urlsafe_b64encode as b64encode, urlsafe_b64decode as b64decode
from Crypto.Cipher import AES

class AES4WEB(object):
    """Create a cryptographically secure, url-safe, base-64
    encoded aes cipher from a given password.
    
    # a 16 byte blocksize is (8*16=128 bit) 128 bit block size
    # a 16 byte iv is good, it should be kept with the string so that
    # the same plaintext isn't the same cipher
    
    # a key of 32 bytes is 256 bits but 128 is minimum
    # (a 128 bit key is 16 characters)
    """

    def __init__(self, key):
        self.AES_BLOCK_SIZE = 16 # bytes = 128 bits
        self.iv = self._msg = None
        
        self.set_key(key)
        
    def _gen_iv(self):
        self.iv = ''.join(chr(random.randint(0, 0xFF)) for i in range(self.AES_BLOCK_SIZE))
        
    def _start_encryptor(self):
        if not self.iv: self._gen_iv()
        self.encryptor = AES.new(self.key, AES.MODE_CBC, self.iv)
        
    def set_key(self, key):
        self.key = hashlib.sha256(key).digest()
        
    def encrypt(self, msg):
        """Deal is as so: 
        first put in the amount of padding and record that at the
        start of the string to encrypt.
        second, make an IV
        third, encrypt the msg
        """
        self._start_encryptor()
        pad = self.AES_BLOCK_SIZE - ( (len(msg)+1) % self.AES_BLOCK_SIZE )

        msg = chr(pad) + msg
        
        msg += pad * chr(random.randint(0,0xFF))
        
        cipher = self.iv
        cipher += self.encryptor.encrypt(msg)
        
        self.cipher = b64encode( cipher )
        
        return self.cipher
    
    def decrypt(self, cipher):
        if not len(cipher):
            return ''
        
        cipher = b64decode(cipher)
        self.iv = cipher[:16]
        self._start_encryptor()
        msg = self.encryptor.decrypt(cipher[16:])
        pad = ord(msg[0])
        msg = msg[1:-pad]
        self._msg = msg
        return msg

if __name__ == '__main__':
    msg = """Deal is as so: 
    first put in the amount of padding and record that at the
    start of the string to encrypt.
    second, make an IV
    third, encrypt the msg
    """
    crypto = AES4WEB('password')
    cipher = crypto.encrypt(msg)
    del crypto # not necessary
    crypto = AES4WEB('password')
    msg2 = crypto.decrypt(cipher)
    assert msg == msg2
    print 'all is good'